{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d19310d-68e2-4652-ae03-4f4ca040c161",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "from typing import List\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import mlflow\n",
    "import numpy as np\n",
    "import numpy.typing as npt\n",
    "import pandas as pd\n",
    "from matplotlib.ticker import MaxNLocator\n",
    "from saai.mcm import MCM\n",
    "import seaborn as sns\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "\n",
    "TRACKING_SERVER_URL = \"http://127.0.0.1:5000\"\n",
    "\n",
    "\n",
    "def rename_metric_columns(df):\n",
    "    df.rename(\n",
    "        columns={\n",
    "            \"run_id\": \"ID\",\n",
    "            \"tags.mlflow.parentRunId\": \"pID\",\n",
    "            \"metrics.saai_acc\": \"saai_acc\",\n",
    "            \"metrics.ssc_acc\": \"ssc_acc\",\n",
    "            \"metrics.skssc_acc\": \"skssc_acc\",\n",
    "            \"metrics.ari_acc\": \"ari_acc\",\n",
    "            \"metrics.fmi_acc\": \"fmi_acc\",\n",
    "            \"metrics.tsxm_acc\": \"tsxm_acc\",\n",
    "            \"metrics.saai\": \"saai\",\n",
    "            \"metrics.ssc\": \"ssc\",\n",
    "            \"metrics.skssc\": \"skssc\",\n",
    "            \"metrics.ari\": \"ari\",\n",
    "            \"metrics.fmi\": \"fmi\",\n",
    "            \"metrics.saai_wins\": \"saai_wins\",\n",
    "            \"metrics.ssc_wins\": \"ssc_wins\",\n",
    "            \"metrics.skssc_wins\": \"skssc_wins\",\n",
    "            \"metrics.fmi_wins\": \"fmi_wins\",\n",
    "            \"metrics.ari_wins\": \"ari_wins\",\n",
    "            \"metrics.tsxm_wins\": \"tsxm_wins\",\n",
    "            \"metrics.tsxm_k\": \"tsxm_k\",\n",
    "            \"metrics.mt\": \"mt\",\n",
    "            \"metrics.p1\": \"p1\",\n",
    "            \"metrics.p2\": \"p2\",\n",
    "            \"metrics.n1\": \"n1\",\n",
    "            \"metrics.corr_avg\": \"corr_avg\",\n",
    "            \"params.data_generator.dim\": \"d\",\n",
    "            \"params.data_generator.r_sync\": \"r_sync\",\n",
    "            \"params.data_generator.n_classes\": \"n_classes\",\n",
    "            \"params.data_generator.lags.1\": \"lags\",\n",
    "            \"tags.run_id\": \"run_id\",\n",
    "            \"params.gp.k\": \"k\",\n",
    "            \"params.gp.sensor\": \"sensor\",\n",
    "        },\n",
    "        inplace=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f50f4baf-de15-4f1d-8f40-83f7c74b8d69",
   "metadata": {},
   "source": [
    "# Synthetic ICS data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "180fbd4d-2c9c-4140-b8d1-0a973a70c893",
   "metadata": {},
   "outputs": [],
   "source": [
    "COLUMNS = [\n",
    "    \"ID\",\n",
    "    \"pID\",\n",
    "    \"saai_acc\",\n",
    "    \"ssc_acc\",\n",
    "    \"skssc_acc\",\n",
    "    \"ari_acc\",\n",
    "    \"fmi_acc\",\n",
    "    \"tsxm_acc\",\n",
    "    \"saai\",\n",
    "    \"ssc\",\n",
    "    \"skssc\",\n",
    "    \"ari\",\n",
    "    \"fmi\",\n",
    "    \"saai_wins\",\n",
    "    \"ssc_wins\",\n",
    "    \"skssc_wins\",\n",
    "    \"fmi_wins\",\n",
    "    \"ari_wins\",\n",
    "    \"tsxm_wins\",\n",
    "    \"tsxm_k\",\n",
    "    \"mt\",\n",
    "    \"p1\",\n",
    "    \"p2\",\n",
    "    \"n1\",\n",
    "    \"corr_avg\",\n",
    "    \"d\",\n",
    "    \"r_sync\",\n",
    "    \"n_classes\",\n",
    "    \"lags\",\n",
    "    \"run_id\",\n",
    "]\n",
    "mlflow.set_tracking_uri(TRACKING_SERVER_URL)\n",
    "df_raw = mlflow.search_runs(\n",
    "    experiment_names=[\"SAAI\"],\n",
    "    filter_string=\"attributes.status = 'FINISHED' AND tags.run_id != 'edeniss2020_ics'\",\n",
    ")\n",
    "rename_metric_columns(df_raw)\n",
    "df = df_raw[COLUMNS]\n",
    "\n",
    "df[\"r_usync\"] = 1 - df[~df[\"r_sync\"].str.contains(\"\\\\[\")][\"r_sync\"].astype(float)\n",
    "df[[\"lags\", \"corr_avg\"]] = df[[\"lags\", \"corr_avg\"]].fillna(0.0)\n",
    "df[\"lags\"] = df[\"lags\"].astype(int)\n",
    "df[\"corr_avg\"] = df[\"corr_avg\"].astype(float)\n",
    "df[\"n_classes\"] = df[~df[\"n_classes\"].str.contains(\"\\\\[\")][\"n_classes\"].astype(int)\n",
    "df[\"d\"] = df[\"d\"].astype(int)\n",
    "df_agr = df[df[\"pID\"].isna()]\n",
    "\n",
    "LABELS = [\"ARI\", \"FMI\", \"SAAI\", \"SSC\", \"X-Means\"]\n",
    "METRICS = [\"ari_acc\", \"fmi_acc\", \"saai_acc\", \"skssc_acc\", \"tsxm_acc\"]\n",
    "df_agr = df_agr.rename(columns={k: v for k, v in zip(METRICS, LABELS)})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14231c6a-b6c8-4aaa-91b5-d10d94d2b78b",
   "metadata": {},
   "source": [
    "## Line plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e587f29f-f037-4dbe-a587-4217355c66e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set(style=\"whitegrid\")\n",
    "FONTSIZE = 22\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "df_classes_inc = df_agr[df_agr[\"run_id\"] == \"increase_k\"].sort_values(\"n_classes\")\n",
    "df_classes_inc.plot(\n",
    "    x=\"n_classes\", y=LABELS, marker=\"x\", ax=ax, fontsize=FONTSIZE, ylim=(0, 1)\n",
    ")\n",
    "ax.legend(labels=LABELS, fontsize=FONTSIZE - 7, loc=\"upper right\")\n",
    "ax.set_ylabel(\"Accuracy\", fontsize=FONTSIZE)\n",
    "ax.xaxis.set_major_locator(MaxNLocator(integer=True))\n",
    "ax.set_xlabel(\"$K$\", fontsize=FONTSIZE)\n",
    "ax.set_ylim = (0, 1)\n",
    "ax.xaxis.grid(False)\n",
    "ax.set_title(\"Increase K\")\n",
    "fig.tight_layout()\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "df_classes_inc = df_agr[\n",
    "    (df_agr[\"run_id\"] == \"increase_psync\") & df_agr[\"run_id\"]\n",
    "].sort_values(\"r_usync\")\n",
    "df_classes_inc.plot(\n",
    "    x=\"r_usync\",\n",
    "    y=LABELS,\n",
    "    marker=\"x\",\n",
    "    ax=ax,\n",
    "    fontsize=FONTSIZE,\n",
    "    legend=False,\n",
    "    ylim=(0, 1),\n",
    ")\n",
    "ax.set_ylabel(\"Accuracy\", fontsize=FONTSIZE)\n",
    "ax.set_xlabel(\"$1- r_{sync}$\", fontsize=FONTSIZE)\n",
    "ax.xaxis.grid(False)\n",
    "ax.set_title(\"Increase $r_{sync}$\")\n",
    "fig.tight_layout()\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "df_classes_inc = df_agr[df_agr[\"run_id\"] == \"increase_d\"].sort_values(\"d\")\n",
    "df_classes_inc.plot(\n",
    "    x=\"d\", y=LABELS, marker=\"x\", ax=ax, fontsize=FONTSIZE, legend=False, ylim=(0, 1)\n",
    ")\n",
    "ax.set_ylabel(\"Accuracy\", fontsize=FONTSIZE)\n",
    "ax.set_xlabel(\"$D$\", fontsize=FONTSIZE)\n",
    "ax.xaxis.grid(False)\n",
    "ax.set_title(\"Increase D\")\n",
    "fig.tight_layout()\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(12, 5))\n",
    "df_lags = df_agr[\n",
    "    (df_agr[\"run_id\"] == \"shifting_correlation\") & (df_agr[\"lags\"] % 60 == 0)\n",
    "].sort_values(\"lags\")\n",
    "df_lags.plot(x=\"lags\", y=LABELS, ax=ax, fontsize=FONTSIZE, marker=\"x\")\n",
    "ax.plot(\n",
    "    df_lags[\"lags\"], df_lags[\"corr_avg\"], color=\"k\", linestyle=\"--\", label=\"Correlation\"\n",
    ")\n",
    "ax.hlines(1 / 19, -720, +720, color=\"k\", linestyle=\"dotted\")\n",
    "ax.legend(labels=LABELS + [\"$\\\\rho$\"], fontsize=FONTSIZE - 7)\n",
    "ax.set_xticks([-720, -540, -360, -180, 0, 180, 360, 540, 720])\n",
    "ax.set_ylabel(\"Accuracy, $\\\\rho$\", fontsize=FONTSIZE)\n",
    "ax.set_xlabel(\"Lag\", fontsize=FONTSIZE)\n",
    "ax.axvspan(-180, 180, color=\"k\", alpha=0.1)\n",
    "ax.xaxis.grid(False)\n",
    "ax.set_title(\"Shifting Correlation\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f2bf34b-a795-4fbd-a9c3-7de22ec58749",
   "metadata": {},
   "source": [
    "## Correlation & Multi Comparions Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "681f88b7-8a91-4d1e-9b08-6f15704d4e92",
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set(style=\"whitegrid\", rc={\"axes.grid\": False})\n",
    "corr_matrix = df_agr[LABELS].corr()\n",
    "saai_row = corr_matrix.loc[\"SAAI\"]\n",
    "saai_row = saai_row.sort_values(ascending=False)\n",
    "fig, ax = plt.subplots(figsize=(5, 4))\n",
    "im = ax.matshow(saai_row.values.reshape(1, -1), cmap=\"Blues\", vmin=0.45)\n",
    "ax.set_xticks(ticks=np.arange(len(LABELS)))\n",
    "ax.set_xticklabels(labels=saai_row.index, fontsize=16, rotation=45, ha=\"left\")\n",
    "ax.set_yticks([0])\n",
    "ax.set_yticklabels(labels=[\"SAAI\"], fontsize=16)\n",
    "ax.tick_params(axis=\"both\", which=\"both\", length=0)\n",
    "\n",
    "for j in range(len(LABELS)):\n",
    "    if j != 0:\n",
    "        text = ax.text(\n",
    "            j,\n",
    "            0,\n",
    "            f\"{saai_row.iloc[j]:.2f}\",\n",
    "            ha=\"center\",\n",
    "            va=\"center\",\n",
    "            color=\"violet\",\n",
    "            fontweight=\"bold\",\n",
    "            fontsize=16,\n",
    "        )\n",
    "    else:\n",
    "        text = ax.text(\n",
    "            j,\n",
    "            0,\n",
    "            \"$\\\\rho$\",\n",
    "            ha=\"center\",\n",
    "            va=\"center\",\n",
    "            color=\"violet\",\n",
    "            fontweight=\"bold\",\n",
    "            fontsize=16,\n",
    "        )\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73fec3bb-dbe4-4657-a8e5-35260c1d6856",
   "metadata": {},
   "outputs": [],
   "source": [
    "MCM.compare(\n",
    "    df_results=df_agr[LABELS],\n",
    "    row_comparates=\"SAAI\",\n",
    "    fig_size=(10, 3),\n",
    "    font_size=14.5,\n",
    "    include_pvalue_legend=False,\n",
    ")\n",
    "\n",
    "\n",
    "# Sweetspot\n",
    "sweetspot_df = df_agr[\n",
    "    ~((df_agr[\"run_id\"] == \"shifting_correlation\") & (df_agr[\"corr_avg\"] < 0.25))\n",
    "    & ~((df_agr[\"run_id\"] == \"increase_psync\") & (df_agr[\"r_usync\"] > 0.8))\n",
    "]\n",
    "MCM.compare(\n",
    "    df_results=sweetspot_df[LABELS],\n",
    "    row_comparates=\"SAAI\",\n",
    "    fig_size=(10, 3),\n",
    "    font_size=14.5,\n",
    "    include_pvalue_legend=False,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6b10bc2-2bb1-493f-8773-16bde516d156",
   "metadata": {},
   "source": [
    "## Ablation Study"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "361b9c86-b17d-47bd-bea2-2691fa698629",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SAAI:\n",
    "\n",
    "    VERSIONS = {\n",
    "        \"original\": 0,\n",
    "        \"only_p1\": 0,\n",
    "        \"only_p2\": 0,\n",
    "    }\n",
    "\n",
    "    def __init__(self, mt, n1, k, mtw):\n",
    "        self.mt = mt\n",
    "        self.n1 = n1\n",
    "        self.k = k\n",
    "        self.mtw = mtw\n",
    "\n",
    "        self.__calculate()\n",
    "\n",
    "    def __calculate(self):\n",
    "        for version in self.VERSIONS.keys():\n",
    "            method = getattr(self, version)\n",
    "            self.VERSIONS[version] = method()\n",
    "\n",
    "    def original(self):\n",
    "        score = self.mtw * self.mt + (1 - self.mtw) * ((self.k - 1 - self.n1) / (self.k))\n",
    "\n",
    "        return score\n",
    "\n",
    "    def only_p1(self):\n",
    "        score = self.mtw * self.mt + (1 - self.mtw) * ((self.k - 1) / (self.k))\n",
    "\n",
    "        return score\n",
    "\n",
    "    def only_p2(self):\n",
    "        score = self.mtw * self.mt + (1 - self.mtw) * ((self.k - self.n1) / (self.k))\n",
    "\n",
    "        return score\n",
    "\n",
    "\n",
    "def recalc_saai(df_abl, df, mtw):\n",
    "    client = mlflow.MlflowClient()\n",
    "    for i, exp in tqdm(df_abl.iterrows()):\n",
    "        if exp[\"pID\"] is not None:\n",
    "            continue\n",
    "        metrics_by_k = {\n",
    "            \"mt\": {k: [] for k in range(2, 20)},\n",
    "            \"n1\": {k: [] for k in range(2, 20)},\n",
    "        }\n",
    "        saai_by_k = {v: {k: [] for k in range(2, 20)} for v in SAAI.VERSIONS.keys()}\n",
    "        num_reps = 0\n",
    "        for j, run in enumerate(df[df[\"pID\"] == exp[\"ID\"]].iterrows()):\n",
    "            num_reps += 1\n",
    "            run = run[1]\n",
    "            for metric in metrics_by_k.keys():\n",
    "                metricitems = client.get_metric_history(run[\"ID\"], metric)\n",
    "                metricitems.sort(key=lambda x: x.timestamp, reverse=False)\n",
    "                for item in metricitems:\n",
    "                    metrics_by_k[metric][item.step].append(item.value)\n",
    "            for k in range(2, 20):\n",
    "                saai = SAAI(\n",
    "                    metrics_by_k[\"mt\"][k][j],\n",
    "                    metrics_by_k[\"n1\"][k][j],\n",
    "                    k,\n",
    "                    mtw,\n",
    "                )\n",
    "                for v in SAAI.VERSIONS.keys():\n",
    "                    saai_by_k[v][k].append(saai.VERSIONS[v])\n",
    "        __dfs = {}.fromkeys(SAAI.VERSIONS.keys())\n",
    "        for v in SAAI.VERSIONS.keys():\n",
    "            __dfs[v] = pd.DataFrame().from_dict(saai_by_k[v])\n",
    "        kfound = {v: 0 for v in SAAI.VERSIONS.keys()}\n",
    "        for j, row in __dfs[\"original\"].iterrows():\n",
    "            k_true = df_abl.loc[i,'n_classes']\n",
    "            for v in SAAI.VERSIONS.keys():\n",
    "                kfound[v] += int(__dfs[v].iloc[j, :].idxmax() == k_true)\n",
    "        for v in SAAI.VERSIONS.keys():\n",
    "            kfound[v] = kfound[v] / num_reps\n",
    "            df_abl.loc[i, f\"SAAI_{v}\"] = kfound[v]\n",
    "    return df_abl\n",
    "\n",
    "\n",
    "df_abl = df[df[\"saai\"].isna()].copy()\n",
    "df_abl = recalc_saai(df_abl, df, 0.5)\n",
    "\n",
    "LABELS_abl = [\"SAAI\", \"$SAAI_{p1}$\", \"$SAAI_{p2}$\"]\n",
    "METRICS_abl = [\"SAAI_original\", \"SAAI_only_p1\", \"SAAI_only_p2\"]\n",
    "\n",
    "df_abl = df_abl.rename(columns={k: v for k, v in zip(METRICS_abl, LABELS_abl)})\n",
    "\n",
    "MCM.compare(\n",
    "    df_results=df_abl[LABELS_abl],\n",
    "    fig_size=(8, 4),\n",
    "    font_size=12,\n",
    "    include_pvalue_legend=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a647d9d1-0756-4fa4-bc07-f711aa92ee2a",
   "metadata": {},
   "source": [
    "## Selecting $\\lambda$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61379bd8-8857-4a99-b277-4ca3b3ae27b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SAAI():\n",
    "    \n",
    "    VERSIONS = {\n",
    "        'original': 0,\n",
    "    }\n",
    "    \n",
    "    def __init__(self, mt, n1, k, l):\n",
    "        self.mt = mt\n",
    "        self.n1 = n1\n",
    "        self.k = k\n",
    "        self.l = l\n",
    "        \n",
    "        self.__calculate()\n",
    "    \n",
    "    def __calculate(self):\n",
    "        for version in self.VERSIONS.keys():\n",
    "            method = getattr(self, version)\n",
    "            self.VERSIONS[version] = method()\n",
    "            \n",
    "    \n",
    "    def original(self):\n",
    "        score = self.l * self.mt + (1 - self.l) * ((self.k - 1 - self.n1) / (self.k))\n",
    "\n",
    "        return score\n",
    "\n",
    "\n",
    "def recalc_saai_by_l(df_lambda, df):\n",
    "    client = mlflow.MlflowClient()\n",
    "    for i, exp in tqdm(df_lambda.iterrows()):\n",
    "        if exp['pID'] is not None:\n",
    "            continue\n",
    "        metrics_by_k = {\n",
    "            'mt': {k: [] for k in range(2, 20)},\n",
    "            'n1': {k: [] for k in range(2, 20)},\n",
    "        }\n",
    "        saai_by_k = {round(l,2): {k: [] for k in range(2, 20)} for l in np.arange(0,1.01,0.05)}\n",
    "        num_reps = 0\n",
    "        for j, run in tqdm(enumerate(df[df['pID'] == exp['ID']].iterrows())):\n",
    "            num_reps += 1\n",
    "            run = run[1]\n",
    "            for metric in metrics_by_k.keys():\n",
    "                metricitems = client.get_metric_history(run['ID'], metric)\n",
    "                metricitems.sort(key=lambda x: x.step, reverse=False)\n",
    "                for item in metricitems:\n",
    "                    metrics_by_k[metric][item.step].append(item.value)\n",
    "            for k in range(2,20):\n",
    "                for l in np.arange(0,1.01,0.05):\n",
    "                    l = round(l,2)\n",
    "                    saai = SAAI(\n",
    "                        metrics_by_k['mt'][k][j],\n",
    "                        metrics_by_k['n1'][k][j],\n",
    "                        k,\n",
    "                        l\n",
    "                    )\n",
    "                    saai_by_k[l][k].append(saai.VERSIONS['original'])\n",
    "        __dfs = {}.fromkeys([round(l,2) for l in np.arange(0,1.01,0.05)])\n",
    "        for l in np.arange(0,1.01,0.05):\n",
    "            l = round(l,2)\n",
    "            __dfs[l] = pd.DataFrame().from_dict(saai_by_k[l])\n",
    "        kfound = {round(l,2): 0 for l in np.arange(0,1.01,0.05)}\n",
    "        for j, row in __dfs[0.0].iterrows():\n",
    "            k_true = df_lambda.loc[i,'n_classes']\n",
    "            for l in np.arange(0,1.01,0.05):\n",
    "                l = round(l,2)\n",
    "                kfound[l] += int(__dfs[l].iloc[j,:].idxmax() == k_true)\n",
    "        for l in np.arange(0,1.01,0.05):\n",
    "            l = round(l,2)\n",
    "            kfound[l] = kfound[l]/num_reps\n",
    "            df_lambda.loc[i, f\"SAAI_{l}\"] = kfound[l]\n",
    "    return df_lambda\n",
    "\n",
    "\n",
    "df_lambda = df[df['saai'].isna()].copy()\n",
    "df_lambda = recalc_saai_by_l(df_lambda, df)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.hlines(y=df_lambda[f'ari_acc'].mean(), xmin=0, xmax=1, color='tab:blue')\n",
    "ax.hlines(y=df_lambda[f'fmi_acc'].mean(), xmin=0, xmax=1, color='tab:orange')\n",
    "for l in np.arange(0,1.01,0.05):\n",
    "    l = round(l,2)\n",
    "    ax.scatter(l, df_lambda[f'SAAI_{l}'].mean(), color='tab:green')\n",
    "    if l == 0.0:\n",
    "        ax.hlines(y=df_lambda[f'skssc_acc'].mean(), xmin=0, xmax=1, color='tab:red')\n",
    "        ax.hlines(y=df_lambda[f'tsxm_acc'].mean(), xmin=0, xmax=1, color='tab:purple')\n",
    "ax.legend(labels=LABELS, fontsize=15, loc=\"lower right\")\n",
    "ax.set_ylabel('Accuracy', fontsize=22)\n",
    "ax.set_xlabel('$\\lambda$', fontsize=22)\n",
    "ax.xaxis.grid(False) \n",
    "ax.yaxis.grid(True) \n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4be4a2bc-302a-4314-81bd-fbc5d65866f8",
   "metadata": {},
   "source": [
    "# Real ICS data (edeniss2020)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "492b26a3-d317-4736-afdf-9c7abca446c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClusterSequencePlot:\n",
    "\n",
    "    @staticmethod\n",
    "    def create(\n",
    "        X: npt.NDArray,\n",
    "        clusters: npt.NDArray[int],\n",
    "        xticklabels: List[str] = None,\n",
    "        fontsize: int = 14,\n",
    "        ax_: npt.NDArray = None,\n",
    "        cmap=None,\n",
    "    ) -> plt.Figure:\n",
    "        k_clusters = clusters.max() + 1\n",
    "        if cmap is None:\n",
    "            cmap = plt.get_cmap(\"tab10\")\n",
    "\n",
    "        if ax_ is None:\n",
    "            fig, ax = plt.subplots(\n",
    "                ncols=k_clusters, figsize=(k_clusters * 2.5, 3), sharey=False\n",
    "            )\n",
    "            if k_clusters == 1:\n",
    "                ax = [ax]\n",
    "        else:\n",
    "            ax = ax_\n",
    "\n",
    "        for yi in range(k_clusters):\n",
    "            for xx in X[clusters == yi]:\n",
    "                ax[yi].plot(xx.ravel(), \"k-\", alpha=0.2)\n",
    "\n",
    "            if xticklabels is not None:\n",
    "                ax[yi].set_xticks(np.arange(len(xticklabels)))\n",
    "                ax[yi].set_xticklabels(\n",
    "                    xticklabels, rotation=45, ha=\"right\", fontsize=fontsize\n",
    "                )\n",
    "\n",
    "            handles, labels = ax[yi].get_legend_handles_labels()\n",
    "\n",
    "        for yi in range(k_clusters):\n",
    "            ax[yi].set_yticks([])\n",
    "            ax[yi].set_yticklabels([])\n",
    "            if xticklabels is None:\n",
    "                ax[yi].set_xlim((0, ax[yi].get_xlim()[1]))\n",
    "                ax[yi].set_xticks(ax[yi].get_xticks())\n",
    "                ax[yi].set_xticklabels(\n",
    "                    ax[yi].get_xticks().astype(int), fontsize=fontsize\n",
    "                )\n",
    "            # ax[yi].patch.set_facecolor(cmap(yi))\n",
    "            ax[yi].patch.set_alpha(0.5)\n",
    "            ax[yi].set_title(\n",
    "                f\"$|C_{{{yi}}}|$={len(X[clusters == yi])}\", fontsize=fontsize\n",
    "            )\n",
    "\n",
    "        if ax_ is None:\n",
    "            fig.tight_layout()\n",
    "            return fig, ax\n",
    "        else:\n",
    "            return ax\n",
    "\n",
    "\n",
    "COLUMNS = [\n",
    "    \"ID\",\n",
    "    \"saai\",\n",
    "    \"ssc\",\n",
    "    \"skssc\",\n",
    "    \"mt\",\n",
    "    \"p1\",\n",
    "    \"p2\",\n",
    "    \"n1\",\n",
    "    \"run_id\",\n",
    "    \"k\",\n",
    "    \"sensor\",\n",
    "]\n",
    "mlflow.set_tracking_uri(TRACKING_SERVER_URL)\n",
    "df_raw = mlflow.search_runs(\n",
    "    experiment_names=[\"[rewi_fe]_SAAI\"],\n",
    "    filter_string=\"attributes.status = 'FINISHED' AND tags.run_id = 'edeniss2020_ics'\",\n",
    ")\n",
    "rename_metric_columns(df_raw)\n",
    "df_ics = df_raw[COLUMNS]\n",
    "df_ics[\"saai\"] = df_ics[\"saai\"].astype(float)\n",
    "df_ics[\"ssc\"] = df_ics[\"ssc\"].astype(float)\n",
    "df_ics[\"skssc\"] = df_ics[\"skssc\"].astype(float)\n",
    "df_ics[\"k\"] = df_ics[\"k\"].astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ff53431-9f77-428b-a365-0da02e00a516",
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.set(style=\"whitegrid\")\n",
    "df = pd.melt(\n",
    "    df_ics[[\"k\", \"saai\", \"ssc\", \"skssc\"]].rename(\n",
    "        columns={\n",
    "            \"saai\": \"SAAI\",\n",
    "            \"ssc\": \"SSC (tslearn)\",\n",
    "            \"skssc\": \"SSC (sklearn)\",\n",
    "            \"k\": \"K\",\n",
    "        }\n",
    "    ),\n",
    "    id_vars=[\"K\"],\n",
    "    var_name=\"metric\",\n",
    "    value_name=\"score\",\n",
    ")\n",
    "\n",
    "fig, ax = plt.subplots(nrows=1, figsize=(8, 4))\n",
    "sns.barplot(x=\"K\", y=\"score\", hue=\"metric\", data=df.sort_values(\"K\"), ax=ax)\n",
    "ax.legend(\n",
    "    title=\"Cluster Measure\", bbox_to_anchor=(0.7, 0.4), loc=\"upper left\", fontsize=12\n",
    ")\n",
    "ax.tick_params(axis=\"x\", labelsize=16)\n",
    "ax.tick_params(axis=\"y\", labelsize=16)\n",
    "ax.set_xlabel(\"K\", fontsize=16)\n",
    "ax.set_ylabel(\"\", fontsize=16)\n",
    "\n",
    "ax.set_title(\"SAAI optimal K by method.\")\n",
    "\n",
    "# Show the plot\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c2f2192-dc8e-46b2-8f3e-058b2d30777d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from saai.features import DenoisedFeatures\n",
    "from saai.utils import seed_rngs\n",
    "\n",
    "df = df_ics\n",
    "seed_rngs(42)\n",
    "\n",
    "titles = {\n",
    "    \"saai\": \"SAAI\",\n",
    "    \"ssc\": \"SSC (tslearn)\",\n",
    "    \"skssc\": \"SSC (sk learn)\",\n",
    "}\n",
    "\n",
    "for m in [\"saai\", \"skssc\"]:\n",
    "    saai_opt = df.loc[df[m].idxmax()]\n",
    "\n",
    "    local_path = mlflow.artifacts.download_artifacts(\n",
    "        run_id=saai_opt[\"ID\"], artifact_path=f\"clusters_k_{saai_opt['k']}.pkl\"\n",
    "    )\n",
    "    with open(local_path, \"rb\") as file:\n",
    "        X, clusters, y, dim_lookup, idx_lookup = pickle.load(file).values()\n",
    "\n",
    "    transform = DenoisedFeatures(w=5)\n",
    "    X_ = transform.transform(X)\n",
    "    X_ = np.array([x.data for x in X_])\n",
    "\n",
    "    fig, ax = ClusterSequencePlot.create(X_, clusters)\n",
    "    for i, ax_ in enumerate(ax):\n",
    "        ax_.set_title(f\"$C_{{{i+1}}}$\", fontsize=32)\n",
    "        ax_.set_xticks([])\n",
    "\n",
    "    fig.tight_layout()\n",
    "    fig.suptitle(titles[m])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
